#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy
import time
# import getpass
from controller_msgs.srv import FileTransCmd, FileTransStatus

class FileTransferClient:
    def __init__(self):
        rospy.init_node('file_transfer_test_client', anonymous=True)
        
        # 等待服务可用
        rospy.loginfo("等待文件传输服务启动...")
        rospy.wait_for_service('file_trans_cmd')
        rospy.wait_for_service('file_trans_status')
        
        self.cmd_client = rospy.ServiceProxy('file_trans_cmd', FileTransCmd)
        self.status_client = rospy.ServiceProxy('file_trans_status', FileTransStatus)
        
        rospy.loginfo("文件传输服务连接成功！")

    def get_user_input(self):
        """获取用户输入"""
        print("\n" + "="*50)
        print("ROS文件传输客户端")
        print("="*50)
        
        # 选择传输方向
        while True:
            direction = input("请选择传输方向 (1-上传, 2-下载): ").strip()
            if direction == "1":
                trans_mode = 1;
                break
            elif direction == "2":
                trans_mode = 2;
                break
            else:
                print("请输入1或2")
        
        # 获取文件路径
        if trans_mode == "upload":
            air_path = input("请输入天空端文件路径: ").strip()
            ground_path = input("请输入地面端文件路径: ").strip()
        else:
            ground_path = input("请输入地面端文件路径: ").strip()
            air_path = input("请输入天空端文件路径: ").strip()
        
        # 获取服务器信息
        ground_user = input("请输入地面端用户名: ").strip()
        ground_host = input("请输入地面端主机地址: ").strip()
        
        # 获取密码（安全输入）
        # password = getpass.getpass("请输入密码/密钥文件: ")
        password = input("请输入密码/密钥文件: ").strip()
        
        return {
            'trans_mode': trans_mode,
            'air_path': air_path,
            'ground_user': ground_user,
            'ground_host': ground_host,
            'ground_path': ground_path,
            'ground_ssh_key': password
        }

    def start_transfer(self, params):
        """启动文件传输"""
        rospy.loginfo("正在启动文件传输...")
        
        try:
            response = self.cmd_client(
                trans_mode=params['trans_mode'],
                air_path=params['air_path'],
                ground_user=params['ground_user'],
                ground_host=params['ground_host'],
                ground_path=params['ground_path'],
                ground_ssh_key=params['ground_ssh_key']
            )
            
            if response.result:
                rospy.loginfo("传输任务已接受，任务ID: %d", response.trans_id)
                return response.trans_id
            else:
                rospy.logerror("传输任务被拒绝: %s", response.message)
                return None
                
        except rospy.ServiceException as e:
            rospy.logerr("服务调用失败: %s", e)
            return None

    def monitor_transfer(self, trans_id):
        """监控传输进度"""
        if not trans_id:
            return False
        
        rospy.loginfo("开始监控传输进度...")
        print("\n传输进度:")
        print("-" * 40)
        
        last_progress = -1
        start_time = time.time()
        
        try:
            while not rospy.is_shutdown():
                # 查询状态
                status_response = self.status_client(trans_id)
                
                # 显示进度
                if status_response.progress != last_progress:
                    current_time = time.time()
                    elapsed = current_time - start_time
                    
                    print(f"进度: {status_response.progress:6.2f}% | "
                          f"状态: {status_response.status:10} | "
                          f"耗时: {elapsed:6.1f}s")
                    
                    last_progress = status_response.progress
                
                # 检查是否完成
                if status_response.status in [2, 3]:
                    print("-" * 40)
                    if status_response.status == 2:
                        rospy.loginfo("文件传输完成！")
                        rospy.loginfo("最终消息: %s", status_response.message)
                    else:
                        rospy.logerr("文件传输失败！")
                        rospy.logerr("错误信息: %s", status_response.message)
                    break
                
                # 每秒查询一次
                time.sleep(1)
                
            return status_response.status == 2
            
        except rospy.ServiceException as e:
            rospy.logerr("状态查询失败: %s", e)
            return False

    def run(self):
        """主运行循环"""
        while not rospy.is_shutdown():
            try:
                # 获取用户输入
                params = self.get_user_input()
                
                # 启动传输
                trans_id = self.start_transfer(params)
                if not trans_id:
                    continue
                
                # 监控进度
                success = self.monitor_transfer(trans_id)
                
                # 询问是否继续
                if success:
                    print("\n✓ 传输成功完成！")
                else:
                    print("\n✗ 传输失败！")
                
                continue_transfer = input("\n是否继续传输其他文件？(y/n): ").strip().lower()
                if continue_transfer not in ['y', 'yes', '是']:
                    break
                    
            except KeyboardInterrupt:
                print("\n用户中断程序")
                break
            except Exception as e:
                rospy.logerr("发生错误: %s", e)
                continue

def main():
    try:
        client = FileTransferClient()
        client.run()
        print("\n感谢使用文件传输客户端！")
        
    except rospy.ROSInterruptException:
        print("ROS中断")
    except Exception as e:
        print(f"程序异常: {e}")

if __name__ == "__main__":
    main()
